/**
 * 
 */
package edu.usc.viterbi.cs.datastructures;

import static com.yahoo.psushant.datastructures.Precondition.*;

import java.util.ArrayList;
import java.util.List;
import java.util.Random;

// TODO the way height is computed is probably inefficient. modify.
/**
 * @author psushant
 * 
 */
public class AvlTree<K extends Comparable<K>, V> implements Dictionary<K, V> {

    private AvlNode<K, V> root;

    @Override
    public void insert(K key, V value) {
        checkNull(key, "key");
        if (root == null) {
            root = new AvlNode<K, V>(key, value, null, null);
            return;
        }
        insert(root, key, value);
        root = rebalance(root);
    }

    private void insert(AvlNode<K, V> node, K key, V value) {
        if (key.compareTo(node.key) < 0) {
            if (node.left == null) {
                node.left = new AvlNode<K, V>(key, value, null, null);
            } else {
                insert(node.left, key, value);
            }
        } else  if (key.compareTo(node.key) > 0) {
            if (node.right == null) {
                node.right = new AvlNode<K, V>(key, value, null, null);
            } else {
                insert(node.right, key, value);
            }
        } else {
            node.value = value;
        }
    }
    
    /** 
     * Performs a tri-node restructuring.  Assumes the nodes are in one
     * of following configurations:
     *
     * <pre>
     *          z=c       z=c        z=a         z=a
     *         /  \      /  \       /  \        /  \
     *       y=b  t4   y=a  t4    t1  y=c     t1  y=b
     *      /  \      /  \           /  \         /  \
     *    x=a  t3    t1 x=b        x=b  t4       t2 x=c
     *   /  \          /  \       /  \             /  \
     *  t1  t2        t2  t3     t2  t3           t3  t4
     * </pre>
     */
    private AvlNode<K, V> rebalance(AvlNode<K, V> node) {
        if (node == null) {
            return null;
        } 
        node.left = rebalance(node.left);
        node.right = rebalance(node.right);
        int balance = height(node.left) - height(node.right);       // optimize height calculation TODO
        AvlNode<K, V> modSubTree = node;
        if (balance < -1 || balance > 1) {    // unbalanced
            if (balance > 0) {
                AvlNode<K, V> child = node.left;
                int childbalance = height(child.left) - height(child.right);
                if (childbalance > 0) {
                    modSubTree = leftleftrotate(node);
                } else {
                    modSubTree = leftrightrotate(node);
                }
            } else {
                AvlNode<K, V> child = node.right;
                int childbalance = height(child.left) - height(child.right);
                if (childbalance > 0) {
                    modSubTree = rightleftrotate(node);
                } else {
                    modSubTree = rightrightrotate(node);
                }
            }
        }
        return modSubTree;
    }
    
    private boolean checkBalance(AvlNode<K, V> node) {
        if (node == null) {
            return true;
        }
        boolean result = checkBalance(node.left) && checkBalance(node.right);
        int balance = height(node.left) - height(node.right);
//        System.out.println("key = " + node.key + ", balance = " + balance);
        return result && !(balance < -1 || balance > 1);
    }

    private AvlNode<K, V> leftleftrotate(AvlNode<K, V> node) {
        AvlNode<K, V> modSubTree = node.left;
        node.left = modSubTree.right;
        modSubTree.right = node;
        return modSubTree;
    }
    
    private AvlNode<K, V> rightrightrotate(AvlNode<K, V> node) {
        AvlNode<K, V> modSubTree = node.right;
        node.right = modSubTree.left;
        modSubTree.left = node;
        return modSubTree;
    }

    private AvlNode<K, V> leftrightrotate(AvlNode<K, V> node) {
        AvlNode<K, V> modSubTree = node.left.right;
        node.left.right = modSubTree.left;
        modSubTree.left = node.left;
        node.left = modSubTree.right;
        modSubTree.right = node;
        return modSubTree;
    }

    private AvlNode<K, V> rightleftrotate(AvlNode<K, V> node) {
        AvlNode<K, V> modSubTree = node.right.left;
        node.right.left = modSubTree.right;
        modSubTree.right = node.right;
        node.right = modSubTree.left;
        modSubTree.left = node;
        return modSubTree;
    }

    // check if we can use node.height
    private int height(AvlNode<K, V> node) {
        if (node == null) {
            return 0;
        }
        return Math.max(height(node.left), height(node.right)) + 1;
    }

    @Override
    public V search(K key) {

        // TODO Auto-generated method stub
        return null;
    }

    @Override
    public void delete(K key) {
        checkNull(key, "key");
        root = delete(root, key);
        rebalance(root);
    }
    
    private AvlNode<K, V> delete(AvlNode<K, V> node, K key) {
        if (node == null) {
            return null;
        }
        if (key.compareTo(node.key) < 0) {
            node.left = delete(node.left, key);
            return node;
        } else if (key.compareTo(node.key) > 0) {
            node.right = delete(node.right, key);
            return node;
        } else {
            if (node.left == null) {
                return node.right;
            }
            if (node.right == null) {
                return node.left;
            }
            AvlNode<K, V> p = min(node.right);
            p.right = deleteMin(node.right);
            p.left = node.left;
            return p;
        }
    }

    private AvlNode<K, V> min(AvlNode<K, V> node) {
        if (node.left == null) {
            return node;
        }
        return min(node.left);
    }
    
    private AvlNode<K, V> deleteMin(AvlNode<K, V> node) {
        if (node.left == null) {
            return node.right;
        }
        node.left = deleteMin(node.left);
        return node;
    }
    
    public static void main(String[] args) {
        AvlTree<Integer, Object> tree = new AvlTree<Integer, Object>();
        Random r = new Random();
        int n = 10000;
        for (int i = 0; i < n; i++) {
            tree.insert(r.nextInt(n / 10), null);
//            System.out.println("****************************");
//            tree.printTree(tree.root);
        }
        
//        Random r = new Random();
//        for (int i = 0; i < 20; i++) {
//            tree.insert( i % 3, null);
//        }
        
        verifyBalance(tree);
        verifyCorrectness(tree);
        
        
        for (int i = 0; i < n; i++) {
            int deleteKey = r.nextInt(n / 100);
            tree.delete(deleteKey);
//            System.out.println("**************************** + delete " + deleteKey);
//            tree.printTree(tree.root);
            verifyBalance(tree);
            verifyCorrectness(tree);
        }
        
//        System.out.println("*********************");
//        tree.printTree(tree.root);
    }

    private static void verifyCorrectness(AvlTree<Integer, Object> tree) {
        List<Integer> l = new ArrayList<Integer>();
        tree.inOrderWalkList(tree.root, l);
        tree.checkSorted(l);
        if (tree.checkSorted(l)) {
            System.out.println("verified BST correctness using inorder walk");
        } else {
            System.out.println("BST incorrect");
        }
    }

    private static void verifyBalance(AvlTree<Integer, Object> tree) {
        if (tree.checkBalance(tree.root)) {
            System.out.println("balanced");
        } else {
            System.out.println("unbalanced");
        }
    }
    
    /**
     * Internal method to print a subtree in sorted order. Inorder traversal on BST gives sorted list
     * @param node
     *            the node that roots the tree.
     */
    private void printTree(AvlNode<K, V> node) {
        if (node != null) {
            printTree(node.left);
            System.out.println(node.key);
            printTree(node.right);
        }
    }
    
    private void inOrderWalkList(AvlNode<K, V> node, List<K> l) {
        if (node != null) {
            inOrderWalkList(node.left, l);
            l.add(node.key);
            inOrderWalkList(node.right, l);
        }
    }
    
    private boolean checkSorted(List<K> l) {
        for (int i = 0; i < l.size() - 1; i++) {
            if(l.get(i).compareTo(l.get(i+ 1)) > 0) {
                return false;
            }
        }
        return true;
    }
    
}
